import random

import torch.nn as nn
from torch.utils.data import DataLoader,Dataset
import torch
import pandas as pd
import glob
import cv2
import matplotlib.pyplot as plt
import albumentations as A
from albumentations.pytorch import ToTensorV2
from Config import Config
img1_transformer = A.Compose([
    A.Resize(height=Config.img_size, width=Config.img_size),
    A.Normalize(max_pixel_value=255.0, p=1.0),
    ToTensorV2(p=1.0),
])
img2_transformer = A.Compose([
    A.Resize(height=Config.img_size, width=Config.img_size),
    A.Normalize(max_pixel_value=255.0, p=1.0),
    ToTensorV2(p=1.0),
])
class Simdataset(Dataset):
    def __init__(self,path,img1_transformer=img1_transformer,img2_transformer=img2_transformer):
        self.path = path
        self.data = pd.read_csv(self.path+"train_data.csv",sep=',',names=["id","image"])[1:]
        self.img1_transformer = img1_transformer
        self.img2_transformer = img2_transformer
        self.image_path = glob.glob(self.path+"images/*.jpg")
        self.len = len(self.image_path)
        self.image_name = self.data["image"]
        self.image_id = self.data["id"]
    def __len__(self):
        return self.len
    def __getitem__(self, index):
        index = index+1
        get_same_class = random.randint(0,1)
        img_1 = self.image_name[index]
        img_1_id = self.image_id[index]
        image1 = cv2.imread(self.path+'images/'+img_1, cv2.COLOR_BGR2RGB)
        augmented1 = self.img1_transformer(image = image1)
        label =get_same_class
        label = torch.tensor(label).unsqueeze(0)
        if get_same_class == 1:

            img_index = self.data[self.data["image"]==img_1].index.to_list()
            total_index_list = self.data[self.data["id"]==img_1_id].index.to_list()
            total_index_list.remove(img_index[0])
            img_2 = self.data["image"][random.choice(total_index_list)]
            image2 =cv2.imread(self.path+'images/'+img_2,cv2.COLOR_BGR2RGB)
            augmented2 = self.img2_transformer(image=image2)
            return augmented1["image"].float(),augmented2["image"].float(),label.float()
        if get_same_class == 0:

            total_index_list = self.data[self.data["id"] != img_1_id].index.to_list()
            #total_index_list.remove(0)
            #print(total_index_list)
            img_2 = self.data["image"][random.choice(total_index_list)]
            image2 = cv2.imread(self.path + 'images/' + img_2, cv2.COLOR_BGR2RGB)
            augmented2 = self.img1_transformer(image=image2)
            return augmented1["image"].float(),augmented2["image"].float(),label.float()




if __name__ == '__main__':
    path = "/home/one/lhbdata/pet/pet_biometric_challenge_2022/train/"
    dataset = Simdataset(path)
    dataloader = DataLoader(dataset,batch_size=1,shuffle=False)
    count=0
    for i ,(img1,img2,label) in enumerate(dataloader):

       # print(img1.shape,img2.shape)
        print(label)
    #imag_list = glob.glob(path+"images/*.jpg")
    #print(len(imag_list))

